import python
import semmle.python.dataflow.TaintTracking
import Taint

from
    Call call, Expr arg, boolean expected_taint, boolean has_taint, string test_res,
    string taint_string
where
    call.getLocation().getFile().getShortName() = "test.py" and
    (
        call.getFunc().(Name).getId() = "ensure_tainted" and
        expected_taint = true
        or
        call.getFunc().(Name).getId() = "ensure_not_tainted" and
        expected_taint = false
    ) and
    arg = call.getAnArg() and
    (
        not exists(TaintedNode tainted | tainted.getAstNode() = arg) and
        taint_string = "<NO TAINT>" and
        has_taint = false
        or
        exists(TaintedNode tainted | tainted.getAstNode() = arg |
            taint_string = tainted.getTaintKind().toString()
        ) and
        has_taint = true
    ) and
    if expected_taint = has_taint then test_res = "ok" else test_res = "failure"
// if expected_taint = has_taint then test_res = "✓" else test_res = "✕"
select arg.getLocation().toString(), call.getScope().(Function).getName(), arg.toString(),
    taint_string, test_res
